from __future__ import division, with_statement

# set up logging system
import logging
log_format = "%(module)s.%(name)s@%(asctime)s : %(message)s"
logging.basicConfig(
    level=logging.DEBUG,
    datefmt="%H:%M",
    format=log_format,
)
logger = logging.getLogger("")

# import numpy
builtin_any = any
builtin_all = all
from numpy import *
import numpy
import numpy.random
import numpy.linalg as linalg
del any, all # completely remove from namespace to avoid confusion

# import scipy
import scipy
import scipy.ndimage as ndimage
from scipy.misc import imread, imsave, imfilter, imresize
import scipy.weave

import pylab

import threading

import itertools
from itertools import *

import time
import os
import sys
import os.path

from collections import defaultdict
import shelve
import subprocess

from tempfile import NamedTemporaryFile

from ImageColor import getrgb

import pyglet
import pyglet.window.key
import pyglet.gl

# ======================================
# define directories

BIG_TEMP_DIRECTORY = "/tmp/vision"
if not os.path.exists(BIG_TEMP_DIRECTORY):
	os.mkdir(BIG_TEMP_DIRECTORY)

TEMP_DIRECTORY = BIG_TEMP_DIRECTORY

FAST_TEMP_DIRECTORY = "/dev/shm/vision"
if not os.path.exists(FAST_TEMP_DIRECTORY):
	os.mkdir(FAST_TEMP_DIRECTORY)

REPORTS_DIRECTORY = "."

db_filename = os.path.join(FAST_TEMP_DIRECTORY, "things.shelve")

# ======================================
# misc system tools

class Expectation(Exception):
	pass

def expect(statement):
    if not statement:
        report = "expectation not met"
        raise Expectation(report)

def expect_equal(a, b, *extra, **kw_extra):
    """Like assert().

    If an assert raise an exception it's usually meant there is something wrong with the called code.
    if expect* raise an exception its is an indication that something is wrong with the caller. e.g. wrong parameters passed.

    """
    if not (a == b):
        report = "%s == %s (%s) %s" % (a,b,extra,kw_extra)
        raise Expectation(report)

class memoize:
    """memoize(fn) - an instance which acts like fn but memoizes its arguments
       Will only work on functions with non-mutable arguments
    """
    def __init__(self, fn):
        self.fn = fn
        self.memo = {}
    def __call__(self, *args):
        if not self.memo.has_key(args):
            self.memo[args] = self.fn(*args)
        return self.memo[args]

def execute(args, blocking=True, cwd=""):
    args = map(str, args)
    c = os.path.abspath(args[0])
    if os.path.exists(c):
        args[0] = c
    if not cwd:
        cwd = os.getcwd()
    #~ print args, cwd, blocking
    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=cwd)
    if blocking:
        return p.stdout.read().strip()

_temp_counter = count()
def temp_filename(extension=".tmp", pre="janto_", directory=TEMP_DIRECTORY):
	assert extension.startswith("."), extension
	n = _temp_counter.next() % 20 #XXX modula hack to avoid filling temp dir
	filename = os.path.join(directory, "%s%d%s" % (pre, n, extension))
	#~ filename = "/tmp/janto_%d_%d%s" % (os.getpid(), n, extension)
	try:
		os.remove(filename)
	except OSError:
		pass
	return filename

def pickle_to_file(filename, data, bin=True):
	with file(filename, "w") as f:
		pickle.dump(data, f, bin)

def pickle_from_file(filename):
	return pickle.load(file(filename))

def prof(function):
    import hotshot, hotshot.stats
    profile_output_filename = os.path.join(TEMP_DIRECTORY, "%s.prof" % function.__name__)

    logger.info("profiling %s" % function)
    prof = hotshot.Profile(profile_output_filename)
    try:
            result = prof.runcall(function)
    finally:
            logger.info("saving profile results")
            prof.close()
    assert os.path.exists(profile_output_filename), profile_output_filename

    logger.info("loading profile results")
    stats = hotshot.stats.load(profile_output_filename)

    logger.info("pickling stats")
    stats_pickle_filename = profile_output_filename+".stats.pickle"
    f = file(stats_pickle_filename, "w")
    #~ pickle.dump(stats, f)
    f.close()

    stats_dump_filename = os.path.join(REPORTS_DIRECTORY, "%s.stats" % function.__name__)
    stats.dump_stats(stats_dump_filename) # for usage by pprofui.py

    logger.info("unpickling stats")
    #~ stats = pickle.load(file(stats_pickle_filename))

    logger.info("parsing stats")
    stats.strip_dirs()
    stats.sort_stats('time', 'calls')
    stats.print_stats()

    return result

# ======================================

class Video(object):
    def __init__(self, filename, fps=None, live=False):
        self.filename = filename
        self.fps = fps if fps is not None else 25 #XXX detect fps from file if not specified and not live
        self.live = live

    def frames(self, do_change=True):
        source = pyglet.media.load(self.filename)
        prev_data = None
        for n in count():
            im = source.get_next_video_frame()
            if im is None:
                break

            # create frame object
            frame = {}
            frame["n"] = n
            frame["the_time"] = time.time() if self.live else (n / self.fps)

            # clean image
            data = pyglet2numpy(im)
            frame["clean"] = data
            frame["display"] = data.copy()

            # change in image
            if do_change:
                change = zeros(data.shape[:2], dtype=uint8)
                if prev_data is not None:
                    X, Y = change.shape
                    assert len(data.shape) == 3, "must be RGB"
                    code = """
                        #include <math.h>

                        for (int x=0; x < X; x++) {
                                for (int y=0; y < Y; y++) {
                                        if (
                                                (fabs(data(x,y,0) - prev_data(x,y,0)) > 10) ||
                                                (fabs(data(x,y,1) - prev_data(x,y,1)) > 10) ||
                                                (fabs(data(x,y,2) - prev_data(x,y,2)) > 10)
                                            )
                                            {
                                                change(x,y) = 100;
                                            } else {
                                                change(x,y) = 0;
                                            }
                                }
                        }
                    """
                    scipy.weave.inline(code,
                        ["change", "data", "prev_data", "X", "Y"],
                        type_converters=scipy.weave.converters.blitz,
                        compiler = 'gcc')
                    change = scipy.ndimage.median_filter(change, size=7)

                frame["change"] = change
                prev_data = data # save for next frame

            yield frame

# ======================================
# image manitpulation tools

def swap_axes(image, copy=False):
    """swaps axes in place"""
    if copy:
        image = image.copy()
    if len(image.shape) == 2:
        image.shape = image.shape[1], image.shape[0]
    elif len(image.shape) == 3:
        image.shape = image.shape[1], image.shape[0], image.shape[2] #XXX wtf?
    if copy:
        return image
    else:
        return None # hopefully leads to an early failure downstream if used wrong

def load_image(filename):
    image = imread(filename)
    if len(image.shape) == 3 and image.shape[-1] == 1:
        image.shape = image.shape[:2] # remove wrapping of inner dimension
    return image

def save_image(filename, image):
    imsave(filename, image)

def pyglet2numpy(im, copy=True):
    """convert pyglet image to numpy"""
    expect(im is not None)
    data = numpy.fromstring(im.get_data(im.format, im.pitch), dtype=numpy.uint8)
    if copy:
        data = data.copy()
    if len(im.format) == 3:
        data.shape = im.width, im.height, len(im.format)
    else:
        data.shape = im.width, im.height
    swap_axes(data)
    return data

def numpy2pyglet(data, copy=True):
    """convert numpy to pyglet image"""
    #~ http://code.astraw.com/projects/motmot/wiki/pygarrayimage
    #~ from pygarrayimage.arrayimage import ArrayInterfaceImage
    #~ return ArrayInterfaceImage(data)

    if copy:
        swap_axes(data)
        if len(data.shape) == 2:
            width, height = data.shape
            bands = 1
            format = "I"
        else:
            width, height, bands = data.shape
            format = "RGB"
        pitch = width * bands
        im = pyglet.image.ImageData(width, height, format, data.tostring(), pitch)
    else:
        raise Exception("not implemented") #XXX
    return im

def grayscale(image):
    expect_equal(3, len(image.shape))
    return image.mean(-1)

class Bbox(object):
    """
    Immutable bbox starting with top left at (10,10) and bottom right at (20,20)::
        >>> Bbox((10,10,20,20))
        Bbox((10, 10, 20, 20))

    or::
        >>> Bbox([(10,10), (20,20)])
        Bbox((10, 10, 20, 20))

    """

    top_left = property(fget=lambda self: self.points[0])
    bottom_right = property(fget=lambda self: self.points[1])
    x = property(fget=lambda self: self.points[0][0])
    y = property(fget=lambda self: self.points[0][1])
    width = property(fget=lambda self: self.points[1][0]-self.points[0][0])
    height = property(fget=lambda self: self.points[1][1]-self.points[0][1])

    def __init__(self, points, image=None, color=None):
        points = asarray(points)
        if points.shape == (2,2):
                points = points
        elif points.shape == (4,):
                points = asarray([points[:2], points[2:]])
        else:
                expect_equal(points.shape, (2,2))

        self.points = points

        if image is not None:
                assert self.is_inside_image(image), (points, image.size)

        self.color = color if color is not None else (0,0,255)

    def draw(self, image):
        import drawing
        drawing.draw_squares(image, [self.points], fill=False, color=getrgb("blue"))
